from dataclasses import dataclass

from typing import List

from . import base
from .constants import *


@dataclass
class DMAParams:
    short: int = 20  # 短周期均线
    long: int = 60  # 长周期均线
    open_long: bool = False  # 开多单标记
    open_short: bool = False  # 开空单标记
    symbol: str = 'ag2106'

    @property
    def period(self):
        return self.long + 1


@dataclass
class DoubleMovingAverage(base.Strategy):
    params: DMAParams = DMAParams()

    def log(self, *args: str) -> None:
        super().log(*args)

    def on_init(self, context: base.Context):
        context.subscribe(self.params.symbol, '60s', count=self.params.period)         # 订阅行情

    def on_bar(self, context: base.Context, bars: List[base.Bar]):
        prices = context.get_data(self.params.symbol, '60s', self.params.period, fields='close')
        # 计算长短周期均线
        short_avg = prices['close'].rolling(self.params.short).mean()
        long_avg = prices['close'].rolling(self.params.long).mean()

        # 查询持仓
        position_long = context.get_account().get_position(symbol=self.params.symbol, side=1)
        position_short = context.get_account().get_position(symbol=self.params.symbol, side=2)

        # 短均线下穿长均线，做空(即当前时间点短均线处于长均线下方，前一时间点短均线处于长均线上方)
        if long_avg.iloc[-2] < short_avg.iloc[-2] and long_avg.iloc[-1] >= short_avg.iloc[-1]:
            # 无多仓情况下，直接开空
            if not position_long:
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Sell,
                                     position_effect=PositionEffect_Open,
                                     order_type=OrderType_Market)
                self.log(self.params.symbol, '以市价单调空仓到仓位')
            # 有多仓情况下，先平多，再开空(开空命令放在on_order_status里面)
            else:
                self.params.open_short = True
                # 以市价平多仓
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Sell,
                                     position_effect=PositionEffect_Close,
                                     order_type=OrderType_Market)
                self.log(self.params.symbol, '以市价单平多仓')
        # 短均线上穿长均线，做多（即当前时间点短均线处于长均线上方，前一时间点短均线处于长均线下方）
        if short_avg[-2] < long_avg[-2] and short_avg[-1] >= long_avg[-1]:
            # 无空仓情况下，直接开多
            if not position_short:
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Buy,
                                     position_effect=PositionEffect_Open,
                                     order_type=OrderType_Market)
                self.log(self.params.symbol, '以市价单调多仓到仓位')
            # 有空仓的情况下，先平空，再开多(开多命令放在on_order_status里面)
            else:
                self.params.open_long = True
                # 以市价平空仓
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Buy,
                                     position_effect=PositionEffect_Close,
                                     order_type=OrderType_Market)
                self.log(self.params.symbol, '以市价单平空仓')

    def on_order_status(self, context: base.Context, order: base.Order):
        # 查看下单后的委托状态
        status = order.status
        # 成交命令的方向
        side = order.side
        # 交易类型
        effect = order.position_effect
        # 当平仓委托全成后，再开仓
        if status == 3:
            # 以市价开空仓，需等到平仓成功无仓位后再开仓
            # 如果无多仓且side=2（说明平多仓成功），开空仓
            if effect == 2 and side == 2 and self.params.open_short:
                self.params.open_short = False
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Sell,
                                     position_effect=PositionEffect_Open,
                                     order_type=OrderType_Market)
                self.log(self.params.symbol, '以市价单调空仓到仓位')
            # 以市价开多仓,需等到平仓成功无仓位后再开仓
            # 如果无空仓且side=1（说明平空仓成功），开多仓
            if effect == 2 and side == 1 and self.params.open_long:
                self.params.open_long = False
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Buy,
                                     position_effect=PositionEffect_Open,
                                     order_type=OrderType_Market)
                self.log(self.params.symbol, '以市价单调多仓到仓位')


@dataclass
class DTParams:
    symbol: str = 'SHFE.rb2010'
    N: int = 5
    k1: float = 0.2
    k2: float = 0.2
    buy_line: float = 0.
    sell_line: float = 0.


@dataclass
class DualTrust(base.Strategy):
    params: DTParams = DTParams()

    def on_init(self, context: base.Context):
        # 获取当前时间
        # time = context.now.strftime('%H:%M:%S')

        # 如果是交易时间段，等到开盘时间确保进入algo()
        context.schedule(schedule_func=self.algo, date_rule='1d', time_rule='09:00:00')
        context.schedule(schedule_func=self.algo, date_rule='1d', time_rule='21:00:00')

    def algo(self, context: base.Context):
        # 取历史数据
        data = context.history_n(symbol=self.params.symbol,
                                 frequency='1d',
                                 end_time=context.now,
                                 fields='symbol,open,high,low,close',
                                 count=self.params.N + 1)
        # 取开盘价
        # 回测模式下，开盘价可以直接用history_n取到
        if context.mode == 2:
            # 获取当天的开盘价
            current_open = data['open'].loc[self.params.N]
            # 去掉当天的实时数据
            data.drop([self.params.N], inplace=True)
        # 如果是实时模式，开盘价需要用current取到
        else:
            # 获取当天的开盘价
            current_open = context.current(self.params.symbol)[0]['open']
        # 计算Dual Thrust 的上下轨
        HH = data['high'].max()
        HC = data['close'].max()
        LC = data['close'].min()
        LL = data['low'].min()
        rng = max(HH - LC, HC - LL)
        self.params.buy_line = current_open + rng * self.params.k1  # 上轨
        self.params.sell_line = current_open - rng * self.params.k2  # 下轨

    def on_bar(self, context: base.Context, bars: List[base.Bar]):
        # 取出订阅的这一分钟的bar
        bar = bars[0]
        # 取出买卖线
        buy_line = self.params.buy_line
        sell_line = self.params.sell_line
        # 获取现有持仓
        position_long = context.get_account().get_position(symbol=self.params.symbol, side=PositionSide_Long)
        position_short = context.get_account().get_position(symbol=self.params.symbol, side=PositionSide_Short)
        # 交易逻辑部分
        # 如果超过range的上界
        if bar.close > buy_line:
            if position_long:  # 已经持有多仓，直接返回
                return
            elif position_short:  # 已经持有空仓，平仓再做多。
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Buy,
                                     order_type=OrderType_Market,
                                     position_effect=PositionEffect_Close)
                self.log('市价单平空仓', self.params.symbol)
                context.order_volume(symbol=self.params.symbol,
                                     volume=1,
                                     side=OrderSide_Buy,
                                     order_type=OrderType_Market,
                                     position_effect=PositionEffect_Open)
                self.log('市价单开多仓', self.params.symbol)
            else:  # 没有持仓时，市价开多。
                context.order_volume(symbol=self.params.symbol, volume=1, side=OrderSide_Buy,
                                     order_type=OrderType_Market, position_effect=PositionEffect_Open)
                self.log('市价单开多仓', self.params.symbol)
        # 如果低于range的下界
        elif bar.close < sell_line:
            if position_long:  # 已经持有多仓， 平多再开空。
                context.order_volume(symbol=self.params.symbol, volume=1, side=OrderSide_Sell,
                                     order_type=OrderType_Market, position_effect=PositionEffect_Close)
                self.log('市价单平多仓', self.params.symbol)
                context.order_volume(symbol=self.params.symbol, volume=1, side=OrderSide_Sell,
                                     order_type=OrderType_Market, position_effect=PositionEffect_Open)
                self.log('市价单开空仓', self.params.symbol)
            elif position_short:  # 已经持有空仓，直接返回。
                return
            else:  # 没有持仓，直接开空
                context.order_volume(symbol=self.params.symbol, volume=1, side=OrderSide_Sell,
                                     order_type=OrderType_Market, position_effect=PositionEffect_Open)
                self.log('市价单开空仓', self.params.symbol)
